# -*-coding:utf-8 -*-

'''
@File       : word_discovery_by_cut.py
@Author     : HW Shen
@Date       : 2020/9/22
@Desc       :  在计算 "凝固度" 时，我们只需要计算 "二字片段" 的凝固度，省掉了更多字片段的凝固度计算
'''

from collections import defaultdict
from tqdm import tqdm
from math import log
import re

import os

# import pymongo
# db = pymongo.MongoClient().baike.items
# def texts():
#     for a in db.find(no_cursor_timeout=True).limit(1000000):
#         yield a['content']


class FindNewWords(object):

    def __init__(self, min_count=10, min_pmi=0):

        self.min_count = min_count
        self.min_pmi = min_pmi
        self.chars, self.pairs = defaultdict(int), defaultdict(int)  # 如果 key不存在，那么就用int()函数
        self.total = 0.  # 初始化一个值，int()的默认结果为 0

    def text_filter(self, texts):
        """ 预切断句子，以免得到太多无意义（不是中文、英文、数字）的字符串 """

        for a in tqdm(texts):
            # 这个正则表达式匹配的是任意非中文、非英文、非数字，因此它的意思就是用任意非中文、非英文、非数字的字符断开句子
            for t in re.split(u'[^\u4e00-\u9fa50-9a-zA-Z]+', a):
                if t: yield t

    def count(self, texts):
        """ 计数函数，计算单字出现频数、相邻两字出现的频数 """

        for text in self.text_filter(texts):
            self.chars[text[0]] += 1
            for i in range(len(text) - 1):
                self.chars[text[i+1]] += 1
                self.pairs[text[i:i+2]] += 1
                self.total += 1

        self.chars = {k: v for k, v in self.chars.items() if v >= self.min_count}  # 最少频数过滤
        self.pairs = {k: v for k, v in self.pairs.items() if v >= self.min_count}  # 最少频数过滤
        self.strong_segments = set()

        # 根据互信息找出比较 “密切” 的邻字
        for k, v in self.pairs.items():
            _ = log(self.total * v / (self.chars[k[0]] * self.chars[k[1]]))
            if _ >= self.min_pmi:
                self.strong_segments.add(k)

    def find_words(self, texts):

        """ 根据前述结果来找词语 """

        self.words = defaultdict(int)
        for text in self.text_filter(texts):
            s = text[0]
            for i in range(len(text) - 1):
                if text[i:i + 2] in self.strong_segments:  # 如果比较 “密切 ” 则不断开
                    s += text[i+1]
                else:
                    self.words[s] += 1  # 否则断开，前述片段作为一个词来统计
                    s = text[i+1]

        self.words = {i: j for i, j in self.words.items() if j >= self.min_count}  # 最后再次根据频数过滤


if __name__ == '__main__':

    texts = [line.strip() for line in open("../data/origin/ftp_1.txt", encoding="utf-8").readlines()]

    # 限定值 min_pmi 设定为 1
    fw = FindNewWords(min_count=16, min_pmi=1)
    # fw.count(texts())
    # fw.find_words(texts())
    fw.count(texts)
    fw.find_words(texts)

    import pandas as pd
    words = pd.Series(fw.words).sort_values(ascending=False)
    print(words)

